#include "UDPGroup.h"
#include <WinSock2.h>
#include "UDPClient.h"
#include "../ErrorLog.h"
#include "../Configuration.h"
#include "../OverlappedData.h"


UDPGroup::UDPGroup()
{

}

UDPGroup::~UDPGroup()
{
	Terminate();
}


unsigned long _stdcall UDPGroup::WorkerThread(void* lpParameter)
{
	UDPGroup* group = reinterpret_cast<UDPGroup*>(lpParameter);
	while (true)
	{
		DWORD numBytes = 0;
		ULONG_PTR completionKey = 0;
		LPOVERLAPPED overlapped = nullptr;
		BOOL result = GetQueuedCompletionStatus(group->completionPort, &numBytes, &completionKey, &overlapped, INFINITE);
		if (overlapped == nullptr) {
			break;
		}
		if (!result) {
			int error = WSAGetLastError();
			OverlappedData* overlappedData = reinterpret_cast<OverlappedData*>(overlapped);
			group->RemoveOnMainThread(overlappedData);
			continue;
		}

		else
		{
			OverlappedData* overlappedData = reinterpret_cast<OverlappedData*>(overlapped);
			if (overlappedData->Type == OVERLAPPED_TYPE_RECV)
			{
				NetParam* _param = new NetParam();
				memcpy(_param->addr, &overlappedData->Address, overlappedData->AddressLen);
				_param->user = overlappedData;
				_param->data = new char[numBytes];
				memcpy(_param->data, overlappedData->Buffer.buf, numBytes);
				_param->size = numBytes;
				group->InvokeAsync([](void* o, void* p) {
					UDPGroup* group = (UDPGroup*)o;
					NetParam* param = (NetParam*)p;
					OverlappedData* data = reinterpret_cast<OverlappedData*>(param->user);
					sockaddr_in* addr = (sockaddr_in*)param->addr;
					if (group->onRecv && data->Peer)
						group->onRecv(param->data, param->size, addr, data->Peer);
					}, group, _param);
				if (!group->BeginReceive(overlappedData))
				{
					group->RemoveOnMainThread(overlappedData);
				}
			}
			else
			{
				group->RemoveOnMainThread(overlappedData);
			}
		}
	}
	return 0;
}

bool UDPGroup::Initialize()
{
	Terminate();
	completionPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, nullptr, 0, 0);
	if (completionPort == nullptr) {
		OutputError(NULL, "Failed to create completion port:  %d", WSAGetLastError());
		return false;
	}
	
	SYSTEM_INFO sysInfo;
	GetSystemInfo(&sysInfo);
	int workHandleCount = sysInfo.dwNumberOfProcessors;
	for (int i = 0; i < workHandleCount; i++) {
		HANDLE hThread = CreateThread(NULL, 0, WorkerThread, this, 0, NULL);
		if (hThread == NULL)
			continue;
		threadHandles.push_back(hThread);
	}
	return true;
}

void UDPGroup::Terminate()
{
	for (const auto& handle : threadHandles) {
		PostQueuedCompletionStatus(completionPort, 0, 0, NULL);
	}

	WaitForMultipleObjects(threadHandles.size(), threadHandles.data(), TRUE, INFINITE);
	for (const auto& handle : threadHandles)
		CloseHandle(handle);
	threadHandles.clear();
	CloseHandle(completionPort);
	completionPort = nullptr;
	while(!clients.empty())
		RemoveClient(*clients.begin());
	while (!overlappedDataSet.empty())
	{
		OverlappedData* data = *overlappedDataSet.begin();
		RemoveContext(data);
		OverlappedData_Free(&data);
	}
}

bool UDPGroup::BeginReceive(OverlappedData* data)
{
	ZeroMemory(&data->Overlapped, sizeof(OVERLAPPED));
	if (WSARecvFrom(data->Socket, &data->Buffer, 1, &data->BytesTransferred, &data->Flags, (sockaddr*)&data->Address, &data->AddressLen, &data->Overlapped, NULL) == SOCKET_ERROR) {
		int error = WSAGetLastError();
		if (error != ERROR_IO_PENDING) {
			return false;
		}
	}
	return true;
}


void UDPGroup::InsertContext(OverlappedData* data)
{
	overlappedDataSetMux.lock();
	auto item = overlappedDataSet.find(data);
	if (item == overlappedDataSet.end())
		overlappedDataSet.insert(data);
	else
		throw;
	overlappedDataSetMux.unlock();
}

void UDPGroup::RemoveContext(OverlappedData* data)
{
	overlappedDataSetMux.lock();
	auto item = overlappedDataSet.find(data);
	if (item != overlappedDataSet.end())
		overlappedDataSet.erase(data);
	else
		throw;
	overlappedDataSetMux.unlock();
}

void UDPGroup::RemoveOnMainThread(OverlappedData* data)
{
	NetParam* _param = new NetParam();
	_param->user = data;
	InvokeAsync([](void* o, void* p) {
		UDPGroup* group = (UDPGroup*)o;
		NetParam* param = (NetParam*)p;
		OverlappedData* data = reinterpret_cast<OverlappedData*>(param->user);
		group->RemoveContext(data);
		OverlappedData_Free(&data);
		}, this, _param);
}

bool UDPGroup::AddClient(UDPClient* client, int port)
{
	auto item = clients.find(client);
	if (item != clients.end())
		return false;
	client->SetCompletionPort(completionPort);
	if (!client->Initialize(port))
		return false;
	OverlappedData* data = reinterpret_cast<OverlappedData*>(client->GetRecvContext());
	InsertContext(data);
	data->Peer = client;
	if (!BeginReceive(data))
		return false;
	client->RegisterInsertCallback(&UDPGroup::InsertContext, this);
	client->RegisterRemoveCallback(&UDPGroup::RemoveContext, this);
	clients.insert(client);
	return true;
}

void UDPGroup::RemoveClient(UDPClient* client)
{
	auto item = clients.find(client);
	if (item == clients.end())
		return;
	client->Terminate();
	clients.erase(item);
	WakeAll();
}